import pandas as pd
import numpy as np
import pandas_ta as ta
from typing import Dict, Any

class TechnicalIndicators:
    """技术指标计算类"""
    
    def __init__(self):
        """初始化技术指标计算器"""
        pass
    
    def calculate_all(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        计算所有技术指标
        
        Args:
            df: 包含OHLCV数据的DataFrame
            
        Returns:
            添加了技术指标的DataFrame
        """
        df = df.copy()
        
        # 移动平均线
        df['ma_5'] = ta.sma(df['close'], length=5)
        df['ma_10'] = ta.sma(df['close'], length=10)
        df['ma_20'] = ta.sma(df['close'], length=20)
        df['ma_50'] = ta.sma(df['close'], length=50)
        df['ma_200'] = ta.sma(df['close'], length=200)
        
        # 指数移动平均线
        df['ema_5'] = ta.ema(df['close'], length=5)
        df['ema_10'] = ta.ema(df['close'], length=10)
        df['ema_20'] = ta.ema(df['close'], length=20)
        
        # MACD
        macd = ta.macd(df['close'])
        df['macd'] = macd['MACD_12_26_9']
        df['macd_signal'] = macd['MACDs_12_26_9']
        df['macd_hist'] = macd['MACDh_12_26_9']
        
        # RSI
        df['rsi_14'] = ta.rsi(df['close'], length=14)
        
        # 布林带
        bollinger = ta.bbands(df['close'])
        df['upper_band'] = bollinger['BBU_20_2.0']
        df['middle_band'] = bollinger['BBM_20_2.0']
        df['lower_band'] = bollinger['BBL_20_2.0']
        
        # 成交量指标
        df['volume_sma_20'] = ta.sma(df['volume'], length=20)
        df['volume_ema_20'] = ta.ema(df['volume'], length=20)
        
        # ATR
        df['atr'] = ta.atr(df['high'], df['low'], df['close'])
        
        # 随机指标
        stoch = ta.stoch(df['high'], df['low'], df['close'])
        df['stoch_k'] = stoch['STOCHk_14_3_3']
        df['stoch_d'] = stoch['STOCHd_14_3_3']
        
        # 趋势指标
        df['adx'] = ta.adx(df['high'], df['low'], df['close'])['ADX_14']
        
        # 动量指标
        df['mom'] = ta.mom(df['close'])
        df['roc'] = ta.roc(df['close'])
        
        # 清理缺失值
        df = df.fillna(method='ffill').fillna(method='bfill')
        
        return df
    
    def calculate_custom_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        计算自定义特征
        
        Args:
            df: 包含技术指标的DataFrame
            
        Returns:
            添加了自定义特征的DataFrame
        """
        df = df.copy()
        
        # 价格动量
        df['price_momentum'] = df['close'].pct_change()
        
        # 波动率
        df['volatility'] = df['close'].rolling(window=20).std()
        
        # 趋势强度
        df['trend_strength'] = abs(df['macd']) / df['volatility']
        
        # 价格位置
        df['price_position'] = (df['close'] - df['lower_band']) / (df['upper_band'] - df['lower_band'])
        
        # 成交量趋势
        df['volume_trend'] = df['volume'] / df['volume_sma_20']
        
        # RSI趋势
        df['rsi_trend'] = df['rsi_14'].diff()
        
        # 清理缺失值
        df = df.fillna(method='ffill').fillna(method='bfill')
        
        return df
